﻿using FaissOperation.FaissIndexFlat;
using FaissOperation.Interface;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json;


namespace FaissOperation.LiteVectorCollection
{
    public class LiteVectorIndex : IVectorCollection<float[]>
    {
        private const int DEFAULT_DIMENSION = 1024; // 默认维度1024


        private readonly List<float[]> _vectors = new();
        private readonly List<int> _ids = new();

        private int _dimension = DEFAULT_DIMENSION;

        public int Dimension
        {
            get { return _dimension; }
            set
            {
                if (_vectors != null && _vectors.Count > 0)
                    throw new InvalidOperationException("在集合中有向量存在的情况下，不允许修改维度。");

                _dimension = value;
            }
        }

        public bool NormalizeOnAddAndQuery { get; set; } = false;

        public MetricType MetricTypeInUse { get; set; } = MetricType.Cosine;

        public void AddVector(int id, float[] vector)
        {
            if (vector.Length != Dimension)
                throw new ArgumentException("向量维度不匹配");

            var vec = (float[])vector.Clone();

            if (NormalizeOnAddAndQuery)
                DistanceProvider.SimpleNormalize(vec);

            RemoveVector(id); // 确保ID唯一性，先删除旧的向量（如果存在）

            _ids.Add(id);
            _vectors.Add(vec);
        }

        public void AddVectors(IEnumerable<(int id, float[] Vector)> items)
        {
            foreach (var (id, vector) in items)
            {
                AddVector(id, vector);
            }
        }

        public void RemoveVector(int id)
        {
            int index = _ids.IndexOf(id);

            if (index >= 0 && index < _vectors.Count)
            {
                _ids.RemoveAt(index);
                _vectors.RemoveAt(index);
            }
        }

        public void RemoveVectors(IEnumerable<int> ids)
        {
            foreach (var id in ids)
            {
                RemoveVector(id);
            }
        }

        public IEnumerable<SearchResult> Search(float[] query, int topK)
        {
            // 维度不匹配，不会返回任何数据。
            if (query.Length != Dimension)
            {
                return [];
            }

            var q = (float[])query.Clone();

            if (NormalizeOnAddAndQuery)
            {
                DistanceProvider.SimpleNormalize(q);
            }

            var results = new List<SearchResult>(_vectors.Count);

            for (int i = 0; i < _vectors.Count; i++)
            {
                float d = DistanceProvider.Distance(q, _vectors[i], MetricTypeInUse);
                results.Add(new SearchResult(_ids[i], d));
            }

            if (MetricTypeInUse == MetricType.L2)
            {
                return results.OrderBy(r => r.distance).Take(topK).ToList();
            }
            else
            {
                return results.OrderByDescending(r => r.distance).Take(topK).ToList();
            }

        }

        public void Clear()
        {
            this._ids.Clear();
            this._vectors.Clear();
        }

        public PlainCollectionData GetData()
        {
            var data = new PlainCollectionData
            {
                MetricType = MetricTypeInUse,
                Dimension = Dimension,
                NormalizeOnAdd = NormalizeOnAddAndQuery,
                Vectors = _vectors,
                Ids = _ids
            };
            return data;
        }

        public void LoadData(PlainCollectionData data)
        {
            this._ids.Clear();
            this._vectors.Clear();

            this.Dimension = data.Dimension;
            this.NormalizeOnAddAndQuery = data.NormalizeOnAdd;
            this.MetricTypeInUse = data.MetricType;

            this._ids.AddRange(data.Ids);
            this._vectors.AddRange(data.Vectors);
        }

        public long Count => _vectors.Count;

        public bool IsTrained => true;

        public void Destory()
        {
            this.Clear();
        }

        public void Train(float[][] vectors)
        {
            throw new NotSupportedException("LiteVectorIndex 不支持训练操作。");
        }
    }


}
